# -*- coding: utf-8 -*-
"""
Created on Wed Sep 25 22:22:00 2019

@author: swls_vanlance
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import sys
import cv2
from tqdm import tqdm
import shutil
import numpy as np
from collections import defaultdict

root_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(root_dir)
from os_tool.os_tool import get_all_filepath_level


def video_to_img(file_path, file_stem):
    cap = cv2.VideoCapture(file_path)
    new_dir = os.path.dirname(file_path) + os.sep + os.path.basename(file_path).split('.')[0]
    if not os.path.exists(new_dir):
        os.makedirs(new_dir)

    count = 0
    go = True
    while 1:
        ret, frame = cap.read()
        # cv2.imshow("capture",frame)
        # cv2.moveWindow('capture',100,100)
        if not ret:
            break
        # if cv2.waitKey(1) & 0xFF == ord('c'):
        # if count % 3 == 0:
        save_path = new_dir + os.sep +file_stem +'_' + os.path.basename(file_path).split(".")[0] + '_{}.jpg'.format(str(count)) 
        # print(save_path)
        cv2.imwrite(save_path, frame[:955, :])
        # print(np.mean(frame))
        count += 1
            
        # if cv2.waitKey(1) & 0xFF == ord('b'):
            # break
    print(count)
    cap.release()
    # cv2.destroyAllWindows()


def get_some_img_from_every_dir(root_dir, dest_dir, image_num):
    all_files = get_all_filepath_level(root, ["jpg"])
    all_files = sorted(all_files, key=lambda x: x[1])

    data_dict = defaultdict(list)
    for p in all_files:
        data_dict[p[1]].append(p[0])

    print(len(data_dict.keys()))

    for p in data_dict.keys():
        files = data_dict[p]

        np.random.shuffle(files)
        np.random.shuffle(files)
        if len(files) < 5:
            copy_files = files
        else:
            copy_files = files[:5]
        for f in copy_files:
            shutil.copy(f, dest_dir + os.sep + os.path.basename(f))

                
if __name__ == "__main__":
    # root = sys.argv[1]
    # all_files = get_all_filepath_level(root, ["AVI"])
    # for f in tqdm(all_files):
    #     video_to_img(f[0], f[1])

    root = sys.argv[1]
    dest_dir = sys.argv[2]
    get_some_img_from_every_dir(root,dest_dir,  5)